package memory

import (
	"context"
	"errors"
	"fmt"
	"github.com/cloudwego/hertz/pkg/common/hlog"
	"time"
)

type Memory[C comparable, T any] struct {
	ctx        context.Context
	cacheKey   string
	cacheTTL   time.Duration
	missDataFn missFn[C, T]
}

type missFn[C comparable, T any] func(ctx context.Context, ids []C) (map[C]T, error)

type Option[C comparable, T any] func(c *Memory[C, T])

func NewMemory[C comparable, T any](ctx context.Context, opts ...Option[C, T]) *Memory[C, T] {
	ins := &Memory[C, T]{
		ctx:      ctx,
		cacheTTL: 5 * time.Minute,
	}

	for _, opt := range opts {
		opt(ins)
	}

	return ins
}

func WithCacheKey[C comparable, T any](key string) Option[C, T] {
	return func(c *Memory[C, T]) {
		c.cacheKey = key
	}
}

func WithMissDataFn[C comparable, T any](fn missFn[C, T]) Option[C, T] {
	return func(c *Memory[C, T]) {
		c.missDataFn = fn
	}
}

func WithCacheTTL[C comparable, T any](ttl time.Duration) Option[C, T] {
	return func(c *Memory[C, T]) {
		c.cacheTTL = ttl
	}
}

func (m Memory[C, T]) MGet(ids ...C) (map[C]T, error) {
	if err := m.checkParams(); err != nil {
		return nil, err
	}
	var dataMap = make(map[C]T)
	var missIDs []C
	for _, id := range ids {
		if v, found := GetClient().Get(fmt.Sprintf(m.cacheKey, id)); found {
			dataMap[id] = v.(T)
		} else {
			missIDs = append(missIDs, id)
		}
	}

	if len(missIDs) > 0 {
		tempMap, err := m.missDataFn(m.ctx, missIDs)
		if err != nil {
			hlog.CtxErrorf(m.ctx, "missDataFn data err: %v", err)
			return nil, err
		}

		for k, v := range tempMap {
			dataMap[k] = v
			GetClient().Set(fmt.Sprintf(m.cacheKey, k), v, m.cacheTTL)
		}
	}

	return dataMap, nil
}

func (m Memory[C, T]) checkParams() (err error) {
	if m.missDataFn == nil {
		return errors.New("miss data function")
	}
	if m.cacheKey == "" {
		return errors.New("cache key is empty")
	}
	return nil
}
